"""
@Time: 2021/1/24 下午 6:56
@Author: jinzhuan
@File: ner_toolkit.py
@Desc: 
"""
import torch
from cognlp import *

tokenize_toolkit = TokenizeToolkit()
words = tokenize_toolkit.run('Ontario is the most populous province in Canada.')

ner_toolkit = NerToolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/ner/trex/toolkit/model.pkl',
    vocabulary_path='../../../cognlp/data/ner/trex/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[5],
    max_seq_length=256)
ner_result = ner_toolkit.run(words)
print(ner_result)
